import jwt
from rest_framework_jwt.authentication import BaseJSONWebTokenAuthentication
from rest_framework_jwt.authentication import jwt_decode_handler
from rest_framework.exceptions import AuthenticationFailed


class JWTAuthentication(BaseJSONWebTokenAuthentication):
    def authenticate(self, token):
        if token is None:
            return None
        try:
            # token => payload
            payload = jwt_decode_handler(token)
        except jwt.ExpiredSignature:
            raise AuthenticationFailed('token已过期')
        except:
            raise AuthenticationFailed('非法用户')
        user = self.authenticate_credentials(payload)
        return user

    # 自定义校验规则：auth token jwt，auth为前盐，jwt为后盐
    def parse_jwt_token(self, jwt_token):
        tokens = jwt_token.split()
        if len(tokens) != 3 or tokens[0].lower() != 'auth' or tokens[2].lower() != 'jwt':
            return None
        return tokens[1]
